# pdf_compressor.py
import fitz
import os
import argparse
from pathlib import Path


def compress_pdf(input_path, output_path, quality=75, dpi=150):
    """
    压缩 PDF 文件

    参数:
        input_path (str): 输入 PDF 文件路径
        output_path (str): 输出 PDF 文件路径
        quality (int): 图片质量 (0-100), 默认为 75
        dpi (int): 图片分辨率, 默认为 150
    """
    try:
        # 打开 PDF
        doc = fitz.open(input_path)
        print(f"正在压缩: {input_path}")

        # 遍历每一页
        for page_num in range(len(doc)):
            page = doc[page_num]
            image_list = page.get_images(full=True)

            # 遍历当前页的所有图片
            for image_index, img in enumerate(image_list, start=1):
                xref = img[0]
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]
                image_ext = base_image["ext"]

                # 仅处理 JPEG 和 PNG 图片
                if image_ext.lower() in ["jpeg", "jpg", "png"]:
                    # 压缩图片
                    pix = fitz.Pixmap(doc, xref)

                    # 如果是 RGB 图片，直接压缩
                    if str(pix.colorspace) == "fitz.csRGB":
                        pix.save(f"temp_image.{image_ext}", quality=quality, dpi=(dpi, dpi))
                    # 如果是 CMYK 图片，转换为 RGB 再压缩
                    else:
                        pix = fitz.Pixmap(fitz.csRGB, pix)
                        pix.save(f"temp_image.{image_ext}", quality=quality, dpi=(dpi, dpi))

                    # 读取压缩后的图片
                    with open(f"temp_image.{image_ext}", "rb") as f:
                        new_image_bytes = f.read()

                    # 替换 PDF 中的图片
                    doc.update_image(xref, image=new_image_bytes)

                    # 删除临时文件
                    os.remove(f"temp_image.{image_ext}")

        # 保存压缩后的 PDF
        doc.save(output_path, garbage=4, deflate=True)
        doc.close()

        # 计算压缩率
        original_size = os.path.getsize(input_path)
        compressed_size = os.path.getsize(output_path)
        compression_ratio = 100 * (1 - compressed_size / original_size)

        print(f"压缩完成: {output_path}")
        print(f"原始大小: {original_size / 1024:.2f} KB")
        print(f"压缩后大小: {compressed_size / 1024:.2f} KB")
        print(f"压缩率: {compression_ratio:.2f}%")

    except Exception as e:
        print(f"压缩失败: {e}")


def main():
    parser = argparse.ArgumentParser(description="PDF 压缩工具")
    parser.add_argument("-i", "--input", required=True, help="输入 PDF 文件路径")
    parser.add_argument("-o", "--output", help="输出 PDF 文件路径")
    parser.add_argument("-q", "--quality", type=int, default=75, help="图片质量 (0-100), 默认 75")
    parser.add_argument("-d", "--dpi", type=int, default=150, help="图片分辨率, 默认 150")

    args = parser.parse_args()

    # 处理输出路径
    if not args.output:
        input_file = Path(args.input)
        args.output = str(input_file.parent / f"{input_file.stem}_compressed{input_file.suffix}")

    # 压缩 PDF
    compress_pdf(args.input, args.output, args.quality, args.dpi)


if __name__ == "__main__":
    main()
